from logging import lastResort

from selenium import webdriver
from time import sleep, time
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains
import threading
import time
from queue import Queue
from selenium.webdriver.common.keys import Keys
import requests
import request
import uuid
import fileUtils
import cvUtils
import cv2
import queryChromeHistory
from PIL import Image
import downloadUtils
import imgUtils
import qiniuUpload
import ShelveDB
import MD5Utils
import nsfw
import re

import os
from selenium.webdriver.chrome.options import Options

imageList = []
db = ShelveDB.ShelveDB("aiModelDownload2")  # 键值对存储对象


def startFetch(url, queue, timePeriod, maxPage=20, localRepeatJump=False, debug=True):
    global imageList

    """开始抓取脚本
    localRepeatJump 本地重复判断不重复提交
    """
    # os.system("taskkill /f /im chrome.exe")

    # 设置下载路径
    download_path = r"H:\download\tmp"
    chrome_history_path = r"H:\download\chrome\profile"

    # 设置 ChromeOptions 来配置下载路径和自动下载
    chrome_options = Options()
    chrome_options.add_experimental_option("prefs", {
        "download.default_directory": download_path,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True,
        "profile.default_content_settings.popups": 0
    })
    # 指定用户数据目录
    chrome_options.add_argument(r"--user-data-dir=H:\download\userdata")
    # 下载文件数据库的保存地址
    chrome_options.add_argument(r"--user-data-dir=" + chrome_history_path)
    chrome_options.add_argument("--start-minimized")  # 设置窗口默认最小化

    # 1.创建Chrome浏览器对象，这会在电脑上在打开一个浏览器窗口
    service = Service('D:\chromedriver-win64\chromedriver.exe')
    driver = webdriver.Chrome(service=service, options=chrome_options)

    # 4.设置浏览器的大小
    driver.set_window_size(1600, 1000)
    # 2.通过浏览器向服务器发送URL请求
    driver.get(url)
    sleep(1)

    # 判断是否需要登录才能访问
    # needLogin = False
    # try:
    #     loginElement = driver.find_element(By.XPATH, "//*[contains(text(), 'Log in to view')]")
    #     print("需要登录才能访问，等待操作登录")
    #     needLogin = True
    #     sleep(1)
    #     return "", "", "需要登录才能访问", "", "", "", "需要登录才能访问"
    # except:
    #     print("无需登录就可以访问")
    #     sleep(1)
    #     needLogin = False

    # 判断是否登录
    # while True:
    #     needLogin2 = False
    #     try:
    #         headerElement = driver.find_element(By.CSS_SELECTOR,
    #                                             "header.mantine-Header-root")
    #         loginElement = headerElement.find_element(By.XPATH, "//*[contains(text(), 'Sign In')]")
    #         print("找到sign In")
    #         needLogin2 = True
    #         sleep(1)
    #     except:
    #         sleep(1)
    #
    #     if (needLogin2 is not True):
    #         try:
    #             stackElement = driver.find_element(By.CSS_SELECTOR,
    #                                                "header.mantine-Stack-root")
    #             loginElement = stackElement.find_element(By.XPATH,
    #                                                      "//*[contains(text(), 'GitHub')]")
    #             print("准备登录github")
    #             needLogin2 = True
    #             sleep(1)
    #         except:
    #             sleep(1)
    #
    #     if (needLogin2 is not True):
    #         try:
    #             stackElement = driver.find_element(By.CSS_SELECTOR,
    #                                                "div.auth-form-body")
    #             loginElement = stackElement.find_element(By.XPATH,
    #                                                      "//*[contains(text(), 'GitHub')]")
    #             print("在github登录页")
    #             needLogin2 = True
    #             sleep(1)
    #         except:
    #             sleep(1)
    #
    #     if needLogin2 is not True:
    #         print("已登录")
    #         break
    #     else:
    #         print("未登录,等待登录后再试")

    sleep(1)

    # 筛选器开关元素
    filtersElement = driver.find_element(By.XPATH,
                                         "//*[contains(text(), 'Filters')]")

    filtersElement.click()

    try:
        sleep(1)
        # 筛选器弹窗元素
        filtersElements = driver.find_element(By.ID,
                                              "mantine-r0-dropdown")

        sleep(1)
        dayElement = filtersElements.find_element(By.XPATH,
                                                  "//*[contains(text(), '" + timePeriod + "')]")
        dayElement.click()
        sleep(1)
        filtersElement.click()
    except Exception as e:
        print(e)

    page = 1
    while True:
        if (page >= maxPage):
            break

        sleep(3)

        try:
            # 查找本页的图片
            imgElements = driver.find_elements(By.CSS_SELECTOR,
                                               "div.MasonryContainer_queries__bS_ak div.justify-center div.overflow-hidden a.mantine-Anchor-root img")
            # imgElements = driver.find_elements(By.CSS_SELECTOR,
            #                                    "div.relative flex overflow-hidden rounded-md border-gray-3 bg-gray-0 shadow-gray-4 ")
        except Exception as e:
            # 查找都报错了则直接跳过
            print(e)
            sleep(1)
            continue

        pageCount = len(imgElements)
        successCount = 0
        for element in imgElements:
            try:

                # 点击后进入图片详情页
                src = element.get_attribute("src")
                imageUrl = getImageOriginUrl(src)
                urlMd5 = MD5Utils.string_to_md5(imageUrl)

                # print(imageUrl)
                if localRepeatJump:  # 需要排除重复
                    isDealUrl = db.read(urlMd5, False)
                    if isDealUrl:
                        # 本图片地址已经处理过了
                        print("第" + str(page) + "页，图片的URL本地已处理过，跳过不再处理")
                        continue

                if (isInImageList(imageUrl) is not True):
                    # 向上两级找父元素
                    elementParent = element.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
                    elementParent = elementParent.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级
                    elementParent = elementParent.find_element(By.XPATH, "./..")  # 使用 "./.." 表示父级

                    # 找到点赞的按钮
                    elementPraiseBtn = elementParent.find_element(By.CSS_SELECTOR,
                                                                  "button.mantine-focus-auto span.mantine-Button-inner span.mantine-Button-label")

                    praiseCount = 0
                    try:
                        praiseCount = extract_and_numbers(elementPraiseBtn.text)
                        print(praiseCount)
                    except:
                        pass

                    # 在父元素内找子元素
                    infoBtnElement = elementParent.find_element(By.CSS_SELECTOR,
                                                                "div.relative.flex.items-center button ")

                    infoBtnElement.click()
                    # sleep(2)
                    prompt = None

                    replyCount = 0  # 重试提取的次数
                    # 等待弹出框加载完毕,再提取描述词
                    while True:
                        try:
                            if (replyCount >= 5):
                                print("第" + str(page) + "页，查找描述词超时，跳过本图")
                                break
                            # 找到这个元素则认为加载完成
                            dialogElement = driver.find_element(By.ID,
                                                                "headlessui-portal-root")
                            if (dialogElement is not None):
                                # # 根据这个元素判断接口请求是否完成
                                # otherMetadataElement = dialogElement.find_element(By.XPATH,
                                #                                                   "//*[contains(text(), 'Other metadata')]")

                                # 抓取prompt
                                promptElement = dialogElement.find_element(By.XPATH,
                                                                           "//*[contains(text(), 'Prompt')]")

                                if (promptElement is not None):
                                    promptContentElement = dialogElement.find_element(
                                        By.CSS_SELECTOR, "div.mantine-Text-root.text-sm")
                                    prompt = promptContentElement.text
                                    print("第" + str(page) + "页，描述词提取成功")
                                    successCount += 1
                                    break
                                else:
                                    replyCount = replyCount + 1
                                    print("第" + str(page) + "页，没有描述词，跳过本图")

                                # if (otherMetadataElement is not None):
                                #     print("弹出框请求成功,开始抓取描述词")

                                # 如果有描述词的话

                            sleep(1)
                        except:
                            replyCount = replyCount + 1

                            try:
                                otherMetadataElement = dialogElement.find_element(By.XPATH,
                                                                                  "//*[contains(text(), 'Other metadata')]")
                                if (otherMetadataElement is not None):
                                    print("第" + str(page) + "页，本图没有描述词，跳过")
                                    break
                            except:
                                pass
                            print("第" + str(page) + "页，未找到描述词，准备重试" + str(replyCount))
                            sleep(1)
                            continue

                    if prompt is not None:
                        queue.put({
                            "imageUrl": imageUrl,
                            "urlMd5": urlMd5,
                            "prompt": prompt,
                            "praiseCount": praiseCount
                        })

                    sleep(1)
                else:
                    print("已处理过本图跳过重复图片")

            except Exception as e:
                print("提取prompt过程中报错，跳过本元素")
                sleep(1)
                pass

        print("第" + str(page) + "页，完成成功数" + str(successCount) + "/" + str(pageCount))

        # 滚动到页面底部
        # 按页面向下键滚动
        if (successCount == 0):
            try:
                # 没有成功的则默认向下滚动翻页
                for i in range(4):  # 滚动3次
                    driver.find_element(By.TAG_NAME, 'body').send_keys(Keys.PAGE_DOWN)
                sleep(1)
            except:
                pass

        page = page + 1

    queue.put('quit')
    sleep(5)
    return


def extract_and_numbers(text):
    # 使用正则表达式找到所有数字
    numbers = re.findall(r'[+-]?\d+(?:\.\d+)?', text)
    # 将所有数字拼接成一个字符串
    concatenated_numbers = ''.join(numbers)
    return concatenated_numbers


def getImageOriginUrl(url):
    """得到url的原始链接地址"""
    start = url[0:url.find("width") - 1]
    widthUrl = url[url.find("width"):len(url)]
    end = widthUrl[widthUrl.find("/"): len(widthUrl)]

    return start + end


def isInImageList(url):
    """是否已经存在图片列表"""
    global imageList
    for u in imageList:
        if u['imageUrl'] == url:
            return True

    return False


# 定义子线程的工作函数
def threadDealImage(queue, saveDir, baseUrl, debug=True):
    global imageList
    while True:
        # 从队列中获取消息
        message = queue.get()
        if message == 'quit':
            # 如果收到退出消息，则退出循环
            print('子线程接收到退出信号，正在退出...')
            break
        else:
            if (isInImageList(message) is not True):
                imageList.append(message)
                print(f'子线程接收到消息: {message}')

                # if debug:
                #     # 调试模式不处理后面操作
                #     continue
                try:
                    imageUrl = message['imageUrl']
                    img, filePath, fileName = downloadUtils.downloadImg(
                        imageUrl=imageUrl,
                        savePath=saveDir,
                        encodeAble=False, openProxy=True)

                    # 如果是png图片的话需要进行转格式压缩
                    if "png" in fileName.lower():
                        newPath = saveDir + fileUtils.getFileNameWithoutSuffix(fileName) + ".jpg"
                        imgUtils.pngToJpg(filePath, newPath)
                    if "webp" in fileName.lower():
                        newPath = saveDir + fileUtils.getFileNameWithoutSuffix(fileName) + ".jpg"
                        imgUtils.webpToJpg(filePath, newPath)

                    # 获取文件大小
                    file_size = os.path.getsize(filePath)

                    if (file_size > 1024 * 1024):  # 大于1m则需要压缩
                        # 需要压缩图片来进行判断
                        img.thumbnail((1920, 1920),
                                      Image.ANTIALIAS)  # resize image with high-quality
                        # judgeImgPath = config.getTxt2ImgDir() + "\\" + str(
                        #     uuid.uuid1()) + "." + fileUtils.getFileNameSuffix(keyName)
                        img.save(filePath)
                        pass

                    score = nsfw.predictImg(filePath)

                    if debug:
                        # 调试模式不处理后面操作
                        continue

                    fileName = str(uuid.uuid4()).replace('-',
                                                         '') + "." + fileUtils.getFileNameSuffix(
                        fileName)

                    qiniuUpload.upload_file(filePath, fileName)
                    # 设置延时删除等服务器拿到文件后再修改延时删除功能
                    qiniuUpload.deleteAfterDays(fileName, 1)

                    # 获取原始图片的宽度和高度
                    original_width, original_height = img.size

                    uploadTask(baseUrl=baseUrl, prompt=message['prompt'], urlMd5=message['urlMd5'],
                               praiseCount=message['praiseCount'],
                               keyName=fileName, originUrl=imageUrl, width=original_width,
                               height=original_height, nsfw=score,
                               )

                except Exception as e:
                    print("下载文件或上传失败，请检查代理是否打开或链接是否正常" + str(e))
                    pass

            # else:
            #     print("已存在此链接", message)

        # 完成任务后，向队列发送信号
        queue.task_done()

    print("总计：", str(len(imageList)))


def uploadTask(baseUrl, prompt, urlMd5, praiseCount, keyName, originUrl, nsfw, width, height):
    data = {"prompt": prompt, "keyName": keyName, "originUrl": originUrl, "width": width,
            "height": height, "nsfw": nsfw, "praiseCount": praiseCount}
    response = request.sendPost(url=baseUrl + "/api/drawtask/createSampleTask", data=data)
    if (response['code'] == 200):
        db.write(urlMd5, True)

    print("上传结果:", str(response['code']) + "   nsfw:", str(nsfw))


if __name__ == '__main__':
    # 筛选的时间范围 Day,Week,Month,Year,All Time
    timePeriod = "Week"
    baseUrl = "http://aiservice.xsdr.top/expression"
    # baseUrl = "http://127.0.0.1:20001/expression"
    maxPage = 20

    debug = False  # 调试模式
    localRepeatJump = False  # 跳过重复（本地记录）

    # saveDir = r"N:\trainData\image\\"
    saveDir = r"F:\trainData\image\\"

    # 创建一个队列
    queue = Queue()

    # 创建子线程
    thread = threading.Thread(target=threadDealImage, args=(queue, saveDir, baseUrl, debug))
    thread.start()

    startFetch(
        url="https://civitai.com/images", queue=queue, timePeriod=timePeriod, maxPage=maxPage,
        localRepeatJump=localRepeatJump)
